package spring.ai_chat.ollama.chatApi;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;
import spring.ai_chat.ollama.chatApi.pojo.AiResponse;
import spring.ai_chat.ollama.chatApi.pojo.ChatMessage;
import spring.ai_chat.ollama.chatApi.pojo.ModuleReq;
import spring.ai_chat.ollama.chatApi.util.ContentMap;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Component
public class OllamaChat {
    //发起post请求
    RestTemplate restTemplate = new RestTemplate();

    String url = "http://localhost:11434/api/chat";

    ObjectMapper objectMapper = new ObjectMapper();

    @Autowired
    ContentMap contentMap;

    /**
     * 放回ai的回答
     *
     * @param id      用户id
     * @param content 问题
     */
    public String getResponse(Integer id, String content) throws JsonProcessingException {
        for (int i = 0; i < 3; i++) {
            try {
                List<ChatMessage> message = getChatMessage(id, content);

                ModuleReq req = new ModuleReq(message);

                String s = objectMapper.writeValueAsString(req);

                log.info("发送消息给模型 {}", s);

                HttpEntity<String> httpReq = getHttpReq(s);
                //处理响应
                ResponseEntity<String> entity = restTemplate.postForEntity(url, httpReq, String.class);

                Map<String, ChatMessage> mapResp = new HashMap<>();

                System.out.println("body = " + entity.getBody());

                System.out.println("map = " + mapResp.toString());

                AiResponse aiResponse = objectMapper.readValue(entity.getBody(), AiResponse.class);

                if (req.getModel().contains("deepseek")) {
                    //删除think的那一段
                    String[] split = aiResponse.getMessage().getContent().split("</think>");
                    System.out.println("think: " + split[0]);
                    System.out.println("content" + split[1]);
                    aiResponse.getMessage().setContent(split[1]);
                }

                System.out.println("module 回答：" + mapResp.get("message"));

                ChatMessage chatMessage = aiResponse.getMessage();

                //把模型放回的结果也储存到上下文
                contentMap.put(id, chatMessage);

                return chatMessage.getContent();
            } catch (Exception e) {
                log.error("出错啦 此时重新获取 {}", e.getMessage(), e);
            }
        }
        return "妹妹脑子要烧了，你等哈哈";
    }

    /**
     * 从这里获取到要给ai的聊天记录
     *
     * @param id
     * @param content
     */
    private List<ChatMessage> getChatMessage(Integer id, String content) {
        ChatMessage message = new ChatMessage("user", content);
        List<ChatMessage> chatMessages = contentMap.get(id);
        chatMessages.add(message);
        return chatMessages;
    }

    /**
     * 这里获取到发送给ai的请求
     *
     * @param content json格式的ai需要的请求体内容
     */
    HttpEntity<String> getHttpReq(String content) {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        return new HttpEntity<String>(content, headers);
    }

}
